package com.mall.util.ratelimit;

import java.util.concurrent.atomic.AtomicReferenceArray;
import java.util.concurrent.locks.ReentrantLock;

public abstract class LeapArray<T> {

    protected long windowLengthInMs;
    protected long sampleCount;
    protected long intervalInMs;

    protected final AtomicReferenceArray<WindowWrap<T>> array;

    private final ReentrantLock updateLock = new ReentrantLock();

    public LeapArray(int sampleCount, int intervalInMs) {
        this.windowLengthInMs = intervalInMs / sampleCount;
        this.intervalInMs = intervalInMs;
        this.sampleCount = sampleCount;

        this.array = new AtomicReferenceArray<>(sampleCount);
    }

    // 根据时间戳获取到当前索引位置
    public int calculateTimeIdx(long timeMillis) {
        long timeId = timeMillis / windowLengthInMs;

        return (int) (timeId % array.length());
    }

    // 根据时间戳获取当前窗口开始时间
    public long calculateWindowStart(long timeMills) {
        return timeMills - timeMills % windowLengthInMs;
    }

    public WindowWrap<T> currentWindow() {
        return currentWindow(TimeUtil.currentTimeMillis());
    }

    protected abstract T newEmptyBucket(long timeMillis);

    protected abstract WindowWrap<T> resetWindowTo(WindowWrap<T> windowWrap, long windowStart);

    // 获取当前窗口
    public WindowWrap<T> currentWindow(long timeMillis) {
        if(timeMillis < 0) {
            return null;
        }

        int idx = calculateTimeIdx(timeMillis);
        long windowStart = calculateWindowStart(timeMillis);

        while (true) {
            WindowWrap<T> old = array.get(idx);

            if(old == null) {
                WindowWrap<T> windowWrap = new WindowWrap<T>(windowLengthInMs, windowStart, newEmptyBucket(timeMillis));
                if(array.compareAndSet(idx, null, windowWrap)) {
                    return windowWrap;
                }else {
                    Thread.yield();
                }
            }else if(windowStart == old.getWindowStart()) {
                return old;
            }else if(windowStart > old.getWindowStart()) {
                if(updateLock.tryLock()) {
                    try {
                        return resetWindowTo(old, windowStart);
                    }finally {
                        updateLock.unlock();
                    }
                }else {
                    Thread.yield();;
                }
            }else if(windowStart < old.getWindowStart()) {
                return new WindowWrap<T>(windowLengthInMs, windowStart, newEmptyBucket(timeMillis));
            }
        }
    }
}
